import pandas as pd
import json
from pyecharts import options as opts
from pyecharts.charts import Pie, Bar, Grid


def get_model_price_distribution(df):
    """
    获取车型价格分布数据
    参数: df - 预加载的销售数据DataFrame
    """
    try:
        # 数据清洗：提取价格区间
        def extract_price(price_str):
            try:
                if pd.isna(price_str) or price_str == '':
                    return 0.0
                # 提取售价区间的最低值
                price_str = str(price_str)
                if '-' in price_str:
                    min_price, max_price = map(float, price_str.split('-'))
                    return (min_price + max_price) / 2  # 取中间值
                else:
                    return float(price_str)
            except (ValueError, IndexError, AttributeError):
                return 0.0

        df['平均售价'] = df['售价(万元)'].apply(extract_price)

        # 过滤无效数据
        df = df[df['平均售价'] > 0]

        # 定义价格区间
        bins = [0, 10, 20, 30, 50, 100, float('inf')]
        labels = ['0-10万', '10-20万', '20-30万', '30-50万', '50-100万', '100万以上']

        df['价格区间'] = pd.cut(df['平均售价'], bins=bins, labels=labels, right=False)

        # 计算各价格区间的销量占比
        price_dist = df.groupby('价格区间')['销量'].sum().reset_index()
        total_sales = price_dist['销量'].sum()
        price_dist['占比'] = (price_dist['销量'] / total_sales * 100).round(2)

        # 创建饼图
        pie = (
            Pie()
            .add(
                "",
                [list(z) for z in zip(price_dist['价格区间'], price_dist['销量'])],
                radius=["30%", "60%"],
                center=["25%", "50%"],
                label_opts=opts.LabelOpts(formatter="{b}: {d}%")
            )
            .set_global_opts(
                title_opts=opts.TitleOpts(title="车型价格区间分布", pos_left="15%"),
                legend_opts=opts.LegendOpts(orient="vertical", pos_top="15%", pos_left="5%")
            )
        )

        # 创建柱状图
        bar = (
            Bar()
            .add_xaxis(price_dist['价格区间'].tolist())
            .add_yaxis("销量", price_dist['销量'].tolist())
            .set_global_opts(
                title_opts=opts.TitleOpts(title="各价格区间销量", pos_right="15%"),
                xaxis_opts=opts.AxisOpts(axislabel_opts=opts.LabelOpts(rotate=-30)),
                yaxis_opts=opts.AxisOpts(
                    axislabel_opts=opts.LabelOpts(formatter="{value} 台")
                )
            )
        )

        # 组合图表
        grid = (
            Grid()
            .add(pie, grid_opts=opts.GridOpts(pos_right="55%"))
            .add(bar, grid_opts=opts.GridOpts(pos_left="55%"))
        )

        # 将图表转换为 JSON 数据
        chart_json = grid.dump_options_with_quotes()

        return {
            'success': True,
            'chart_data': chart_json,
            'price_distribution': price_dist.to_dict('records')
        }

    except Exception as e:
        return {
            'success': False,
            'error': f'获取车型价格分布数据失败: {str(e)}'
        }